# weights/VGG16_epoch28_params.pth

from PIL import Image
from torch import nn
from models.model import VGG16
from dataset.read_yolo_dataset import ReadYolo
from augmentation.data_augment import DataAugment
from torch.utils.data import DataLoader
from utils.collate import colle

import torchvision
import numpy as np
import cv2
import torch
import argparse

parser = argparse.ArgumentParser(description='VGG16 Testing')
parser.add_argument("--weight_dir",
                    default='weights/selsected_params_epoch51.pth',
                    # default='weights/VGG16_epoch260_params.pth',
                    help="参数路径")
parser.add_argument("--test_dir",
                    default="./dataset/test/img/image_0222.jpg",
                    help="测试图片路径")

args = parser.parse_args()

# 加载模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

net = VGG16(mode="test")
net = net.to(device=device)

# 加载模型参数
net.load_state_dict(torch.load(args.weight_dir))


# 进行推理
def run():
    net.eval()
    # 此时参数不更新，pool与dropout失效
    image = cv2.imread(args.test_dir, cv2.IMREAD_COLOR)
    trans = torchvision.transforms.Compose([
        torchvision.transforms.ToPILImage(),
        torchvision.transforms.Resize(size=(224, 224)),
        torchvision.transforms.ToTensor()
    ])
    image = trans(image).unsqueeze(0).to(device)
    result = torch.argmax(net(image).ravel())
    # result = net(image)
    return result


def run1():
    batch_size = 1
    net.eval()  # 评估模式
    dataset = ReadYolo(trans=DataAugment(), phase="test", device=device)
    # collate_fn 改变数据的传入方式，要求输入必须是tensor
    data = iter(
        DataLoader(dataset, batch_size, drop_last=False, collate_fn=colle))
    pred_list = []
    label_list = []
    for batch, (imgs, targets) in enumerate(data):
        # print(imgs.shape) # (batch_size, 3,224,224)
        pred = net(imgs)
        pred = torch.argmax(pred.ravel())
        pred_list.append(pred.item())  # item只能取标量 bth = 1
        label_list.append(int(targets.reshape(1).item()))
    # 将tensor转换成标量后再进行np.array运算，得出准确率
    pred_list = np.array(pred_list)
    label_list = np.array(label_list)
    temp_list = (pred_list == label_list)

    print(pred_list)
    print(label_list)

    accuracy = (np.sum(temp_list != 0) + 0.0) / len(temp_list)
    print("准确率为{:.3f}".format(accuracy))


if __name__ == '__main__':
    torch.cuda.empty_cache()
    # 部分测试
    # result = run()
    # print(result)

    # 测试整体
    run1()